Skip to content

[FEA] Truly Decoupled Double Buffering for SG KMeans - #2484

Closed
tarang-jain wants to merge 10 commits into
NVIDIA:mainfrom
tarang-jain:prefetch-sg
Closed

[FEA] Truly Decoupled Double Buffering for SG KMeans#2484
tarang-jain wants to merge 10 commits into
NVIDIA:mainfrom
tarang-jain:prefetch-sg

Conversation

@tarang-jain

@tarang-jain tarang-jain commented Aug 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor
Screenshot 2026-08-18 at 4 44 05 PM Profile: H100 + Intel(R) Xeon(R) Platinum 8480CL

Dataset mem type = pinned memory. This speeds up training by a lot. nsys gave throughput for H2D transfers around 50 GB/s. This is a 114 GB dataset. The screenshot above shows the overlap between the two streams with prefetch.

Metric Baseline (2M batch) Prefetch (2M batch) Baseline (4M batch)
profile_mode baseline prefetch baseline
dataset base_falcon_1024_30M.fbin base_falcon_1024_30M.fbin base_falcon_1024_30M.fbin
rows 30,000,000 30,000,000 30,000,000
dimensions 1,024 1,024 1,024
clusters 2,048 2,048 2,048
batch_rows 2,000,000 2,000,000 4,000,000
requested_iterations 20 20 20
pinned_load_seconds 66.3182 59.2441 60.1955
initialization_seconds 0.0169335 0.0170245 0.0172878
fit_seconds 101.064 59.8278 101.737
fit speedup vs. 4M 1.007x 1.700x 1.000x
e2e_seconds 168.422 120.27 162.921
e2e speedup vs. 4M 0.967x 1.355x 1.000x
iterations 20 20 20
inertia 2.14326e+07 2.14325e+07 2.14325e+07

Also the use of pinned memory shows > 50 GB/s of throughput for H2D transfers. This gets much lower with regular host mem.
Screenshot 2026-08-18 at 5 03 57 PM

@copy-pr-bot

copy-pr-bot Bot commented Aug 18, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@tarang-jain tarang-jain self-assigned this Aug 18, 2026
@tarang-jain tarang-jain added non-breaking Introduces a non-breaking change improvement Improves an existing functionality labels Aug 18, 2026
@tarang-jain
tarang-jain marked this pull request as ready for review August 18, 2026 23:25
@tarang-jain
tarang-jain requested a review from a team as a code owner August 18, 2026 23:25
viclafargue
viclafargue reviewed Aug 20, 2026
View reviewed changes

@viclafargue viclafargue left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work, the reported speedup looks promising! To make the changes easier to review and validate, could we keep this PR focused on double-buffered prefetching? The workspace-resource selection, final-inertia optimization, and cluster-cost refactoring could be handled in separate follow-up PRs, allowing each change to be tested and benchmarked independently.

Comment thread cpp/src/cluster/detail/kmeans.cuh Outdated
viclafargue
viclafargue reviewed Aug 20, 2026
View reviewed changes
Comment thread cpp/src/cluster/detail/kmeans.cuh
@tarang-jain

tarang-jain commented Aug 20, 2026

Copy link
Copy Markdown
Contributor Author

The use of the detail namespace API is part of this optimization. The public API currently synchronizes the stream, so it is not allowing concurrent prefetch during inertia computation.

@tarang-jain
tarang-jain requested review from a team as code owners August 20, 2026 15:47
@tarang-jain tarang-jain mentioned this pull request Aug 20, 2026
Open
@tarang-jain

tarang-jain commented Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

Closing in favor of @viclafargue's work here: #2538
That PR has some extra features to recycle the loader at the end of all the batches, so loading continues for the next iteration.

@tarang-jain tarang-jain closed this Sep 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@viclafargue viclafargue viclafargue left review comments

Assignees

@tarang-jain tarang-jain

Labels

improvement Improves an existing functionality non-breaking Introduces a non-breaking change

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@tarang-jain @viclafargue